# -*- coding: utf-8 -*-

"""剑指 Offer II 010. 和为 k 的子数组
给定一个整数数组和一个整数 k ，请找到该数组中和为 k 的连续子数组的个数。

示例 1 :
输入:nums = [1,1,1], k = 2
输出: 2
解释: 此题 [1,1] 与 [1,1] 为两种不同的情况

示例 2 :
输入:nums = [1,2,3], k = 3
输出: 2

提示:
1 <= nums.length <= 2 * 10^4
-1000 <= nums[i] <= 1000
-10^7 <= k <= 10^7"""


class Solution:
    def subarraySum_timeout(self, nums, k: int) -> int:
        """仅仅用前缀和，减少计算，遍历还是超时，见过程可知，T(n)=O(n^2)"""
        sums = [0]
        for n in nums:
            sums.append(n+sums[-1])
        
        count = 0
        end = len(sums)
        for i in range(1, end):
            step_k = sums[i-1]+k
            for j in range(i, end):
                if step_k == sums[j]:
                    count += 1
        return count
    
    def subarraySum(self, nums, k):
        """前缀和往前用，计算到当前前缀和 total，想要得到和为 k 的子组，那么只要前面的前缀和满足 total-k 的要求，两个前缀和之间的子组就是有效的，T(n)=O(n)"""
        count = 0
        # 为了应对一种情况：某个位置的前缀和刚好为k
        sums_table = {0: 1}
        total = 0
        for num in nums:
            total += num

            count += sums_table.get(total-k, 0)
            
            if total in sums_table.keys():
                sums_table[total] += 1
            else:
                sums_table[total] = 1
        
        return count


if __name__ == '__main__':
    so = Solution()
    print(so.subarraySum([1,1,1,1], 2))
    print(so.subarraySum([1,2,3], 3))
    print(so.subarraySum([1], 0))
